{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5df2f69d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchtext\n",
    "\n",
    "\n",
    "from torchtext.data import Field, BucketIterator, Iterator\n",
    "\n",
    "\n",
    "from torchtext import data\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "\n",
    "import spacy\n",
    "# nlp = spacy.load(\"en_core_web_sm\")\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import random\n",
    "import math\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4c472048",
   "metadata": {},
   "outputs": [],
   "source": [
    "# f = open(\"english_python_data.txt\", \"r\")\n",
    "f = open(\"english_python_data.txt\", \"r\", encoding=\"utf-8\")\n",
    "\n",
    "file_lines = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7560f24d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['# write a python program to add two numbers \\n',\n",
       " 'num1 = 1.5\\n',\n",
       " 'num2 = 6.3\\n',\n",
       " 'sum = num1 + num2\\n',\n",
       " \"print(f'Sum: {sum}')\\n\",\n",
       " '\\n',\n",
       " '\\n',\n",
       " '# write a python function to add two user provided numbers and return the sum\\n',\n",
       " 'def add_two_numbers(num1, num2):\\n',\n",
       " '    sum = num1 + num2\\n',\n",
       " '    return sum\\n',\n",
       " '\\n',\n",
       " '\\n',\n",
       " '# write a program to find and print the largest among three numbers\\n',\n",
       " '\\n',\n",
       " 'num1 = 10\\n',\n",
       " 'num2 = 12\\n',\n",
       " 'num3 = 14\\n',\n",
       " 'if (num1 >= num2) and (num1 >= num3):\\n',\n",
       " '   largest = num1\\n']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_lines[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0fe7e971",
   "metadata": {},
   "outputs": [],
   "source": [
    "dps = []\n",
    "dp = None\n",
    "for line in file_lines:\n",
    "  if line[0] == \"#\":\n",
    "    if dp:\n",
    "      dp['solution'] = ''.join(dp['solution'])\n",
    "      dps.append(dp)\n",
    "    dp = {\"question\": None, \"solution\": []}\n",
    "    dp['question'] = line[1:]\n",
    "  else:\n",
    "    dp[\"solution\"].append(line)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1b06b494",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Question no:  1\n",
      "write a python program to add two numbers \n",
      "\n",
      "num1 = 1.5\n",
      "num2 = 6.3\n",
      "sum = num1 + num2\n",
      "print(f'Sum: {sum}')\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  2\n",
      "write a python function to add two user provided numbers and return the sum\n",
      "\n",
      "def add_two_numbers(num1, num2):\n",
      "    sum = num1 + num2\n",
      "    return sum\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  3\n",
      "write a program to find and print the largest among three numbers\n",
      "\n",
      "\n",
      "num1 = 10\n",
      "num2 = 12\n",
      "num3 = 14\n",
      "if (num1 >= num2) and (num1 >= num3):\n",
      "   largest = num1\n",
      "elif (num2 >= num1) and (num2 >= num3):\n",
      "   largest = num2\n",
      "else:\n",
      "   largest = num3\n",
      "print(f'largest:{largest}')\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  4\n",
      "write a program to find and print the smallest among three numbers\n",
      "\n",
      "num1 = 10\n",
      "num2 = 12\n",
      "num3 = 14\n",
      "if (num1 <= num2) and (num1 <= num3):\n",
      "   smallest = num1\n",
      "elif (num2 <= num1) and (num2 <= num3):\n",
      "   smallest = num2\n",
      "else:\n",
      "   smallest = num3\n",
      "print(f'smallest:{smallest}')\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  5\n",
      "Write a python function to merge two given lists into one\n",
      "\n",
      "def merge_lists(l1, l2):\n",
      "    return l1 + l2\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  6\n",
      "Write a program to check whether a number is prime or not\n",
      "\n",
      "num = 337\n",
      "\n",
      "if num > 1:\n",
      "   for i in range(2, num//2 + 1):\n",
      "       if (num % i) == 0:\n",
      "           print(num,\"is not a prime number\")\n",
      "           print(f\"{i} times {num//i} is {num}\")\n",
      "           break\n",
      "   else:\n",
      "       print(f\"{num} is a prime number\")\n",
      "\n",
      "else:\n",
      "   print(f\"{num} is not a prime number\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  7\n",
      "Write a python function that prints the factors of a given number\n",
      "\n",
      "def print_factors(x):\n",
      "   print(f\"The factors of {x} are:\")\n",
      "   for i in range(1, x + 1):\n",
      "       if x % i == 0:\n",
      "           print(i)\n",
      "\n",
      "\n",
      "\n",
      " Question no:  8\n",
      "Write a program to find the factorial of a number\n",
      "\n",
      "num = 13\n",
      "factorial = 1\n",
      "\n",
      "if num < 0:\n",
      "   print(\"No factorials for negative numbers!\")\n",
      "\n",
      "elif num == 0:\n",
      "   print(\"The factorial of 0 is 1\")\n",
      "\n",
      "else:\n",
      "   for i in range(1,num + 1):\n",
      "       factorial = factorial*i\n",
      "   print(f\"The factorial of {num} is {factorial}\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  9\n",
      "Write a python function to print whether a number is negative, positive or zero\n",
      "\n",
      "def check_pnz(num):\n",
      "    if num > 0:\n",
      "       print(\"Positive number\")\n",
      "\n",
      "    elif num == 0:\n",
      "       print(\"Zero\")\n",
      "\n",
      "    else:\n",
      "       print(\"Negative number\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  10\n",
      "Write a program to print the multiplication table of a given number\n",
      "\n",
      "\n",
      "num = 9\n",
      "for i in range(1, 11):\n",
      "   print(f\"{num} x {i} = {num*i}\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  11\n",
      "Write a python function to print powers of 2, for given number of terms\n",
      "\n",
      "def two_power(terms):\n",
      "    result = list(map(lambda x: 2 ** x, range(terms)))\n",
      "\n",
      "    print(f\"The total terms are: {terms}\")\n",
      "    for i in range(terms):\n",
      "       print(f\"2^{i} = {result[i]}\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  12\n",
      "Write a program to filter the numbers in a list which are divisible by a given number\n",
      "\n",
      "my_list = [11, 45, 74, 89, 132, 239, 721, 21]\n",
      "\n",
      "num = 3\n",
      "result = list(filter(lambda x: (x % num == 0), my_list))\n",
      "\n",
      "print(f\"Numbers divisible by {num} are {result}\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  13\n",
      "Write a python function that returns the sum of n natural numbers\n",
      "\n",
      "def sum_natural(num):\n",
      "    if num < 0:\n",
      "       print(\"Please enter a positive number!\")\n",
      "    else:\n",
      "       sum = 0\n",
      "       while(num > 0):\n",
      "           sum += num\n",
      "           num -= 1\n",
      "       return num\n",
      "\n",
      "\n",
      "\n",
      " Question no:  14\n",
      "Write a program to swap first and last elements in a list\n",
      "\n",
      "my_list = [1, 2, 3, 4, 5, 6]\n",
      "my_list[0], my_list[-1] = my_list[-1], my_list[0]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  15\n",
      "Write a python function to find the area of a circle, whose radius is given\n",
      "\n",
      "def findArea(r): \n",
      "    PI = 3.142\n",
      "    return PI * (r*r)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  16\n",
      "Write a program to print the sum of squares of first n natural numbers\n",
      "\n",
      "n = 21\n",
      "sum_n = 0\n",
      "for i in range(1, n+1):\n",
      "    sum_n += i**2\n",
      "print(sum_n)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  17\n",
      "Write a program to print the length of a list\n",
      "\n",
      "my_list = [1, 2, 3, 4, 5, 6, 7, 8, 9]\n",
      "\n",
      "print(len(my_list))\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  18\n",
      "Write a pythno function to print the length of a given tuple\n",
      "\n",
      "my_tuple = (1, 2, 3, 4, 5, 6, 7, 8)\n",
      "\n",
      "print(len(my_tuple))\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  19\n",
      "Write a python function to print the elements of a given list, one element in a line\n",
      "\n",
      "def custom_print(l):\n",
      "    for _ in l:\n",
      "        print(_)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  20\n",
      "Write a python function to remove all the odd numbers from a list and return the remaining list\n",
      "\n",
      "\n",
      "def remove_odd(my_list):\n",
      "    result = list(filter(lambda x: (x % 2 == 0), my_list))\n",
      "    return result\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  21\n",
      "Write a python function to remove all the even numbers from a list and return the remaining list\n",
      "\n",
      "\n",
      "def remove_even(my_list):\n",
      "    result = list(filter(lambda x: (x % 2 != 0), my_list))\n",
      "    return result\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  22\n",
      "Write a function that takes two lists as input and returns a zipped list of corresponding elements\n",
      "\n",
      "\n",
      "def zip_list(list1, list2):\n",
      "    return list(zip(list1, list2))\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  23\n",
      "Write a program to to print the contents of a given file\n",
      "\n",
      "file_name = 'temp.txt'\n",
      "with open(file_name, 'r') as f:\n",
      "    print(f.read())\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  24\n",
      "Write a functin that returns the LCM of two input numbers\n",
      "\n",
      "\n",
      "def lcm(a, b):\n",
      "    if a>b:\n",
      "        min_ = a\n",
      "    else:\n",
      "        min_ = b\n",
      "    while True:\n",
      "        if min_%a==0 and min_%b==0:\n",
      "            break\n",
      "        min_+=1\n",
      "    return min_\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  25\n",
      "Write a program to print the unique elements in a list\n",
      "\n",
      "my_list = [1, 2, 4, 5, 2, 3, 1, 5, 4, 7, 8, 2, 4, 5, 2, 7, 3]\n",
      "\n",
      "print(set(my_list))\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  26\n",
      "Write a function that returns the sum of digits of a given number\n",
      "\n",
      "def digisum(num):\n",
      "    sum_=0\n",
      "    while num > 0:\n",
      "        dig = num % 10\n",
      "        sum_+=dig\n",
      "        num//=10\n",
      "    return sum_\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  27\n",
      "Write a program to check and print whether a number is palindrome or not\n",
      "\n",
      "\n",
      "num = 12321\n",
      "temp = num\n",
      "rev = 0\n",
      "while num > 0:\n",
      "    dig = num % 10\n",
      "    rev = rev*10 + dig\n",
      "    num//=10\n",
      "if temp==rev :\n",
      "    print(\"The number is a palindrome!\")\n",
      "else:\n",
      "    print(\"The number isn't a palindrome!\")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  28\n",
      "Write a function that prints a given value, n number of times\n",
      "\n",
      "def print_n(val, n):\n",
      "    for _ in range(n):\n",
      "        print(val)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  29\n",
      "Write a function to find the area of sqaure\n",
      "\n",
      "def square_area(a):\n",
      "    return a*a\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  30\n",
      "Write a function to find the perimeter of a square\n",
      "\n",
      "def square_perimeter(a):\n",
      "    return 4*a\n",
      "\n",
      "\n",
      "\n",
      " Question no:  31\n",
      "Write a function to find the area of rectangle\n",
      "\n",
      "def rectangle_area(l, b):\n",
      "    return l*b\n",
      "\n",
      "\n",
      "\n",
      " Question no:  32\n",
      "Write a function to find the permieter of a rectangle\n",
      "\n",
      "def rectangle_perimeter(l, b):\n",
      "    return 2*(l+b)\n",
      "\n",
      "\n",
      "\n",
      " Question no:  33\n",
      "Write a python function to find the area of a circle, whose radius is given\n",
      "\n",
      "def findArea(r): \n",
      "    PI = 3.142\n",
      "    return PI * (r*r)\n",
      "\n",
      "\n",
      "\n",
      " Question no:  34\n",
      "Write a function to calculate and return electricity bill. Units used are given. Price per unit is fixed and is increased after 750 units.\n",
      "\n",
      "\n",
      "def calc_elect_bill(units):\n",
      "    if units > 0:\n",
      "        if units <= 750:\n",
      "            return 5*units\n",
      "        else:\n",
      "            return 5*(750) + 7*(units-750)\n",
      "\n",
      "    else:\n",
      "        return -1\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  35\n",
      "Write a function to return day of a week, given the number\n",
      "\n",
      "def give_day(n):\n",
      "    day_dict = {1: 'Sunday', 2: 'Monday', 3: 'Tuesday', 4: 'Wednesday', 5: 'Thursday', 6: 'Friday', 7: 'Saturday'}\n",
      "    return day_dict[n]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  36\n",
      "Write a program to calculate and print the volume of a cylender\n",
      "\n",
      "r = 3\n",
      "h = 5\n",
      "pi = 3.14\n",
      "volume = pi*(r**2)*h\n",
      "print(volume)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  37\n",
      "Write a function to calculate and return the average of input numbers\n",
      "\n",
      "\n",
      "def calc_avg(*args):\n",
      "    if len(args) > 0:\n",
      "        return sum(args)/len(args)\n",
      "    return None\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  38\n",
      "Write a function to calculate compound interest, given p, r, t\n",
      "\n",
      "def comp_int(p, r, t):\n",
      "    amount = p * (1 + (r/100))**t\n",
      "    interest = amount - p\n",
      "    return interest\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  39\n",
      "Write a function to calculate simple interest, given p, r, t\n",
      "\n",
      "def simp_int(p, r, t):\n",
      "    interest = (p*r*t)/100\n",
      "    return interest\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  40\n",
      "Write a program to print a given string, replacing all the vowels with '_'\n",
      "\n",
      "\n",
      "st = \"Where is this going? Could you please help me understand!\"\n",
      "vowels = \"AEIOUaeiou\"\n",
      "\n",
      "for v in vowels:\n",
      "    st = st.replace(v, '_')\n",
      "\n",
      "print(st)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  41\n",
      "Write a functio to check whether a number if perfect or not\n",
      "\n",
      "def is_perfect(n):\n",
      "    sum_ = 0\n",
      "    for i in range(1, n//2 + 1):\n",
      "        if n%i == 0:\n",
      "            sum_+=i\n",
      "    if sum_ == n:\n",
      "        return True\n",
      "    return False\n",
      "\n",
      "\n",
      "\n",
      " Question no:  42\n",
      "Write a function that returns seperate lists of positive and negative numbers from an input list\n",
      "\n",
      "def seperate_pn(l):\n",
      "    pos_list = []\n",
      "    neg_list = []\n",
      "    for _ in l:\n",
      "        if _<0:\n",
      "            neg_list.append(_)\n",
      "        else:\n",
      "            pos_list.append(_)\n",
      "    return pos_list, neg_list\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  43\n",
      "Write a program to find and print the area of a triangle, whose hight and width are given.\n",
      "\n",
      "\n",
      "h = 12\n",
      "w = 11\n",
      "area = 0.5*h*w\n",
      "print(area)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Question no:  44\n",
      "Write a function to find acceleration, given u, v and t\n",
      "\n",
      "\n",
      "def acc(u, v, t):\n",
      "    return (v-u)/t\n",
      "\n",
      "\n",
      "\n",
      " Question no:  45\n",
      "Write a lambda function to multiply two numbers\n",
      "\n",
      "\n",
      "multiply = lambda a, b: a*b\n",
      "\n",
      "\n",
      "\n",
      " Question no:  46\n",
      "Write a lambda function to add two numbers\n",
      "\n",
      "\n",
      "add = lambda a, b: a+b\n",
      "\n",
      "\n",
      "\n",
      " Question no:  47\n",
      "Write a lambda function that gives True if the input number is even otherwise False\n",
      "\n",
      "\n",
      "even = lambda a: True if a%2 == 0 else False\n",
      "\n",
      "\n",
      "\n",
      " Question no:  48\n",
      "Write a lambda function to to give character grom it's ascii value\n",
      "\n",
      "\n",
      "ascii = lambda a: chr(a)\n",
      "\n",
      "\n",
      "\n",
      " Question no:  49\n",
      "Write a lambda function to that gives the number of digits in a number\n",
      "\n",
      "\n",
      "dig_cnt = lambda a: len(str(a))\n",
      "\n",
      "\n",
      "\n",
      " Question no:  50\n",
      "Write a program to to check if a triangle is valid or not, given it's all three angles\n",
      "\n",
      "\n",
      "def is_valid_triangle_angle(a, b c):\n",
      "    if a+b+c == 180:\n",
      "        return True\n",
      "    return False\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "i=0\n",
    "for dp in dps:\n",
    "  print(\"\\n Question no: \", i+1)\n",
    "  i+=1\n",
    "  print(dp['question'][1:])\n",
    "  print(dp['solution'])\n",
    "  if i>49:\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0d6620d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset size: 4957\n"
     ]
    }
   ],
   "source": [
    "print(\"Dataset size:\", len(dps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "20afd76a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenize import tokenize, untokenize\n",
    "import io\n",
    "\n",
    "\n",
    "def tokenize_python_code(python_code_str):\n",
    "    python_tokens = list(tokenize(io.BytesIO(python_code_str.encode('utf-8')).readline))\n",
    "    tokenized_output = []\n",
    "    for i in range(0, len(python_tokens)):\n",
    "        tokenized_output.append((python_tokens[i].type, python_tokens[i].string))\n",
    "    return tokenized_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e9a5b85b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(63, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (54, '('), (1, 'num1'), (54, ','), (1, 'num2'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'sum'), (54, '='), (1, 'num1'), (54, '+'), (1, 'num2'), (4, '\\n'), (1, 'return'), (1, 'sum'), (4, '\\n'), (62, '\\n'), (62, '\\n'), (6, ''), (0, '')]\n"
     ]
    }
   ],
   "source": [
    "tokenized_sample = tokenize_python_code(dps[1]['solution'])\n",
    "print(tokenized_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b7c5342d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def add_two_numbers (num1 ,num2 ):\n",
      "    sum =num1 +num2 \n",
      "    return sum \n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(untokenize(tokenized_sample).decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9f9b7a55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['False', 'None', 'True', 'and', 'as', 'assert', 'async', 'await', 'break', 'class', 'continue', 'def', 'del', 'elif', 'else', 'except', 'finally', 'for', 'from', 'global', 'if', 'import', 'in', 'is', 'lambda', 'nonlocal', 'not', 'or', 'pass', 'raise', 'return', 'try', 'while', 'with', 'yield']\n"
     ]
    }
   ],
   "source": [
    "import keyword\n",
    "\n",
    "print(keyword.kwlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "813fc1f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_tokenize_python_code(python_code_str, mask_factor=0.3):\n",
    "\n",
    "\n",
    "    var_dict = {} # Dictionary that stores masked variables\n",
    "\n",
    "    # certain reserved words that should not be treated as normal variables and\n",
    "    # hence need to be skipped from our variable mask augmentations\n",
    "    skip_list = ['range', 'enumerate', 'print', 'ord', 'int', 'float', 'zip',\n",
    "                 'char', 'list', 'dict', 'tuple', 'set', 'len', 'sum', 'min', 'max']\n",
    "    skip_list.extend(keyword.kwlist)\n",
    "\n",
    "    var_counter = 1\n",
    "    python_tokens = list(tokenize(io.BytesIO(python_code_str.encode('utf-8')).readline))\n",
    "    tokenized_output = []\n",
    "\n",
    "    for i in range(0, len(python_tokens)):\n",
    "      if python_tokens[i].type == 1 and python_tokens[i].string not in skip_list:\n",
    "        \n",
    "        if i>0 and python_tokens[i-1].string in ['def', '.', 'import', 'raise', 'except', 'class']: # avoid masking modules, functions and error literals\n",
    "          skip_list.append(python_tokens[i].string)\n",
    "          tokenized_output.append((python_tokens[i].type, python_tokens[i].string))\n",
    "        elif python_tokens[i].string in var_dict:  # if variable is already masked\n",
    "          tokenized_output.append((python_tokens[i].type, var_dict[python_tokens[i].string]))\n",
    "        elif random.uniform(0, 1) > 1-mask_factor: # randomly mask variables\n",
    "          var_dict[python_tokens[i].string] = 'var_' + str(var_counter)\n",
    "          var_counter+=1\n",
    "          tokenized_output.append((python_tokens[i].type, var_dict[python_tokens[i].string]))\n",
    "        else:\n",
    "          skip_list.append(python_tokens[i].string)\n",
    "          tokenized_output.append((python_tokens[i].type, python_tokens[i].string))\n",
    "      \n",
    "      else:\n",
    "        tokenized_output.append((python_tokens[i].type, python_tokens[i].string))\n",
    "    \n",
    "    return tokenized_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1c96d000",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(63, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (54, '('), (1, 'num1'), (54, ','), (1, 'var_1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'sum'), (54, '='), (1, 'num1'), (54, '+'), (1, 'var_1'), (4, '\\n'), (1, 'return'), (1, 'sum'), (4, '\\n'), (62, '\\n'), (62, '\\n'), (6, ''), (0, '')]\n"
     ]
    }
   ],
   "source": [
    "tokenized_sample = augment_tokenize_python_code(dps[1]['solution'])\n",
    "print(tokenized_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "242944d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def add_two_numbers (num1 ,var_1 ):\n",
      "    sum =num1 +var_1 \n",
      "    return sum \n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(untokenize(tokenized_sample).decode('utf-8'))\n",
    "     \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "20c06e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "python_problems_df = pd.DataFrame(dps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3820a7ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>solution</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>write a python program to add two numbers \\n</td>\n",
       "      <td>num1 = 1.5\\nnum2 = 6.3\\nsum = num1 + num2\\npri...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>write a python function to add two user provi...</td>\n",
       "      <td>def add_two_numbers(num1, num2):\\n    sum = nu...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>write a program to find and print the largest...</td>\n",
       "      <td>\\nnum1 = 10\\nnum2 = 12\\nnum3 = 14\\nif (num1 &gt;=...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>write a program to find and print the smalles...</td>\n",
       "      <td>num1 = 10\\nnum2 = 12\\nnum3 = 14\\nif (num1 &lt;= n...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Write a python function to merge two given li...</td>\n",
       "      <td>def merge_lists(l1, l2):\\n    return l1 + l2\\n...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0       write a python program to add two numbers \\n   \n",
       "1   write a python function to add two user provi...   \n",
       "2   write a program to find and print the largest...   \n",
       "3   write a program to find and print the smalles...   \n",
       "4   Write a python function to merge two given li...   \n",
       "\n",
       "                                            solution  \n",
       "0  num1 = 1.5\\nnum2 = 6.3\\nsum = num1 + num2\\npri...  \n",
       "1  def add_two_numbers(num1, num2):\\n    sum = nu...  \n",
       "2  \\nnum1 = 10\\nnum2 = 12\\nnum3 = 14\\nif (num1 >=...  \n",
       "3  num1 = 10\\nnum2 = 12\\nnum3 = 14\\nif (num1 <= n...  \n",
       "4  def merge_lists(l1, l2):\\n    return l1 + l2\\n...  "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python_problems_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d332711f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4957, 2)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python_problems_df.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e937cb2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed(0)\n",
    "msk = np.random.rand(len(python_problems_df)) < 0.85 # Splitting data into 85% train and 15% validation\n",
    "\n",
    "train_df = python_problems_df[msk]\n",
    "val_df = python_problems_df[~msk]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a427b345",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4211, 2)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.shape\n",
    "     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ac5b36f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(746, 2)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "\n",
    "val_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1e687767",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "14a91093",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "\n",
    "spacy_en = spacy.load('en_core_web_sm')\n",
    "\n",
    "Input = data.Field(tokenize = lambda text: [tok.text for tok in spacy_en.tokenizer(text)],\n",
    "            init_token='', \n",
    "            eos_token='', \n",
    "            lower=True)\n",
    "\n",
    "Output = data.Field(tokenize = augment_tokenize_python_code,\n",
    "                    init_token='', \n",
    "                    eos_token='', \n",
    "                    lower=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2362fe25",
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [('Input', Input),('Output', Output)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "213f6398",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_example = []\n",
    "val_example = []\n",
    "\n",
    "train_expansion_factor = 100\n",
    "for j in range(train_expansion_factor):\n",
    "  for i in range(train_df.shape[0]):\n",
    "      try:\n",
    "          ex = data.Example.fromlist([train_df.question[i], train_df.solution[i]], fields)\n",
    "          train_example.append(ex)\n",
    "      except:\n",
    "          pass\n",
    "\n",
    "for i in range(val_df.shape[0]):\n",
    "    try:\n",
    "        ex = data.Example.fromlist([val_df.question[i], val_df.solution[i]], fields)\n",
    "        val_example.append(ex)\n",
    "    except:\n",
    "        pass "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4282f594",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[' ', 'write', 'a', 'python', 'program', 'to', 'add', 'two', 'numbers']\n",
      "[(63, 'utf-8'), (1, 'var_1'), (54, '='), (2, '1.5'), (4, '\\n'), (1, 'num2'), (54, '='), (2, '6.3'), (4, '\\n'), (1, 'sum'), (54, '='), (1, 'var_1'), (54, '+'), (1, 'num2'), (4, '\\n'), (1, 'print'), (54, '('), (3, \"f'Sum: {sum}'\"), (54, ')'), (4, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'python', 'function', 'to', 'add', 'two', 'user', 'provided', 'numbers', 'and', 'return', 'the', 'sum']\n",
      "[(63, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (54, '('), (1, 'num1'), (54, ','), (1, 'var_1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'sum'), (54, '='), (1, 'num1'), (54, '+'), (1, 'var_1'), (4, '\\n'), (1, 'return'), (1, 'sum'), (4, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'program', 'to', 'find', 'and', 'print', 'the', 'largest', 'among', 'three', 'numbers']\n",
      "[(63, 'utf-8'), (62, '\\n'), (1, 'var_1'), (54, '='), (2, '10'), (4, '\\n'), (1, 'num2'), (54, '='), (2, '12'), (4, '\\n'), (1, 'num3'), (54, '='), (2, '14'), (4, '\\n'), (1, 'if'), (54, '('), (1, 'var_1'), (54, '>='), (1, 'num2'), (54, ')'), (1, 'and'), (54, '('), (1, 'var_1'), (54, '>='), (1, 'num3'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'largest'), (54, '='), (1, 'var_1'), (4, '\\n'), (6, ''), (1, 'elif'), (54, '('), (1, 'num2'), (54, '>='), (1, 'var_1'), (54, ')'), (1, 'and'), (54, '('), (1, 'num2'), (54, '>='), (1, 'num3'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'largest'), (54, '='), (1, 'num2'), (4, '\\n'), (6, ''), (1, 'else'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'largest'), (54, '='), (1, 'num3'), (4, '\\n'), (6, ''), (1, 'print'), (54, '('), (3, \"f'largest:{largest}'\"), (54, ')'), (4, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'program', 'to', 'find', 'and', 'print', 'the', 'smallest', 'among', 'three', 'numbers']\n",
      "[(63, 'utf-8'), (1, 'var_1'), (54, '='), (2, '10'), (4, '\\n'), (1, 'num2'), (54, '='), (2, '12'), (4, '\\n'), (1, 'num3'), (54, '='), (2, '14'), (4, '\\n'), (1, 'if'), (54, '('), (1, 'var_1'), (54, '<='), (1, 'num2'), (54, ')'), (1, 'and'), (54, '('), (1, 'var_1'), (54, '<='), (1, 'num3'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'var_2'), (54, '='), (1, 'var_1'), (4, '\\n'), (6, ''), (1, 'elif'), (54, '('), (1, 'num2'), (54, '<='), (1, 'var_1'), (54, ')'), (1, 'and'), (54, '('), (1, 'num2'), (54, '<='), (1, 'num3'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'var_2'), (54, '='), (1, 'num2'), (4, '\\n'), (6, ''), (1, 'else'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'var_2'), (54, '='), (1, 'num3'), (4, '\\n'), (6, ''), (1, 'print'), (54, '('), (3, \"f'smallest:{smallest}'\"), (54, ')'), (4, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'python', 'function', 'to', 'merge', 'two', 'given', 'lists', 'into', 'one']\n",
      "[(63, 'utf-8'), (1, 'def'), (1, 'merge_lists'), (54, '('), (1, 'l1'), (54, ','), (1, 'l2'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'return'), (1, 'l1'), (54, '+'), (1, 'l2'), (4, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'program', 'to', 'check', 'whether', 'a', 'number', 'is', 'prime', 'or', 'not']\n",
      "[(63, 'utf-8'), (1, 'num'), (54, '='), (2, '337'), (4, '\\n'), (62, '\\n'), (1, 'if'), (1, 'num'), (54, '>'), (2, '1'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'for'), (1, 'i'), (1, 'in'), (1, 'range'), (54, '('), (2, '2'), (54, ','), (1, 'num'), (54, '//'), (2, '2'), (54, '+'), (2, '1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '       '), (1, 'if'), (54, '('), (1, 'num'), (54, '%'), (1, 'i'), (54, ')'), (54, '=='), (2, '0'), (54, ':'), (4, '\\n'), (5, '           '), (1, 'print'), (54, '('), (1, 'num'), (54, ','), (3, '\"is not a prime number\"'), (54, ')'), (4, '\\n'), (1, 'print'), (54, '('), (3, 'f\"{i} times {num//i} is {num}\"'), (54, ')'), (4, '\\n'), (1, 'break'), (4, '\\n'), (6, ''), (6, ''), (1, 'else'), (54, ':'), (4, '\\n'), (5, '       '), (1, 'print'), (54, '('), (3, 'f\"{num} is a prime number\"'), (54, ')'), (4, '\\n'), (62, '\\n'), (6, ''), (6, ''), (1, 'else'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'print'), (54, '('), (3, 'f\"{num} is not a prime number\"'), (54, ')'), (4, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'python', 'function', 'that', 'prints', 'the', 'factors', 'of', 'a', 'given', 'number']\n",
      "[(63, 'utf-8'), (1, 'def'), (1, 'print_factors'), (54, '('), (1, 'x'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'print'), (54, '('), (3, 'f\"The factors of {x} are:\"'), (54, ')'), (4, '\\n'), (1, 'for'), (1, 'i'), (1, 'in'), (1, 'range'), (54, '('), (2, '1'), (54, ','), (1, 'x'), (54, '+'), (2, '1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '       '), (1, 'if'), (1, 'x'), (54, '%'), (1, 'i'), (54, '=='), (2, '0'), (54, ':'), (4, '\\n'), (5, '           '), (1, 'print'), (54, '('), (1, 'i'), (54, ')'), (4, ''), (6, ''), (6, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'program', 'to', 'print', 'the', 'multiplication', 'table', 'of', 'a', 'given', 'number']\n",
      "[(63, 'utf-8'), (62, '\\n'), (1, 'num'), (54, '='), (2, '9'), (4, '\\n'), (1, 'for'), (1, 'i'), (1, 'in'), (1, 'range'), (54, '('), (2, '1'), (54, ','), (2, '11'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'print'), (54, '('), (3, 'f\"{num} x {i} = {num*i}\"'), (54, ')'), (4, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'python', 'function', 'to', 'print', 'powers', 'of', '2', ',', 'for', 'given', 'number', 'of', 'terms']\n",
      "[(63, 'utf-8'), (1, 'def'), (1, 'two_power'), (54, '('), (1, 'var_1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'result'), (54, '='), (1, 'list'), (54, '('), (1, 'map'), (54, '('), (1, 'lambda'), (1, 'x'), (54, ':'), (2, '2'), (54, '**'), (1, 'x'), (54, ','), (1, 'range'), (54, '('), (1, 'var_1'), (54, ')'), (54, ')'), (54, ')'), (4, '\\n'), (62, '\\n'), (1, 'print'), (54, '('), (3, 'f\"The total terms are: {terms}\"'), (54, ')'), (4, '\\n'), (1, 'for'), (1, 'i'), (1, 'in'), (1, 'range'), (54, '('), (1, 'var_1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '       '), (1, 'print'), (54, '('), (3, 'f\"2^{i} = {result[i]}\"'), (54, ')'), (4, ''), (6, ''), (6, ''), (0, '')]\n",
      "\n",
      "[' ', 'write', 'a', 'program', 'to', 'filter', 'the', 'numbers', 'in', 'a', 'list', 'which', 'are', 'divisible', 'by', 'a', 'given', 'number']\n",
      "[(63, 'utf-8'), (1, 'my_list'), (54, '='), (54, '['), (2, '11'), (54, ','), (2, '45'), (54, ','), (2, '74'), (54, ','), (2, '89'), (54, ','), (2, '132'), (54, ','), (2, '239'), (54, ','), (2, '721'), (54, ','), (2, '21'), (54, ']'), (4, '\\n'), (62, '\\n'), (1, 'num'), (54, '='), (2, '3'), (4, '\\n'), (1, 'result'), (54, '='), (1, 'list'), (54, '('), (1, 'filter'), (54, '('), (1, 'lambda'), (1, 'x'), (54, ':'), (54, '('), (1, 'x'), (54, '%'), (1, 'num'), (54, '=='), (2, '0'), (54, ')'), (54, ','), (1, 'my_list'), (54, ')'), (54, ')'), (4, '\\n'), (62, '\\n'), (1, 'print'), (54, '('), (3, 'f\"Numbers divisible by {num} are {result}\"'), (54, ')'), (4, ''), (0, '')]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for ex in train_example[:10]:\n",
    "    print(ex.Input)\n",
    "    print(ex.Output)\n",
    "    print()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e892f8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = data.Dataset(train_example, fields)\n",
    "valid_data =  data.Dataset(val_example, fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cb243a83",
   "metadata": {},
   "outputs": [],
   "source": [
    "Input.build_vocab(train_data, min_freq = 0)\n",
    "Output.build_vocab(train_data, min_freq = 0)\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "7e844b85",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torchtext.vocab.Vocab at 0x21867d070d0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Output.vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2fde8028",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_vocab(vocab, path):\n",
    "    import pickle\n",
    "    output = open(path, 'wb')\n",
    "    pickle.dump(vocab, output)\n",
    "    output.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c1dea787",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "eab8efc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(63, 'utf-8'),\n",
       " (1, 'var_1'),\n",
       " (54, '='),\n",
       " (2, '1.5'),\n",
       " (4, '\\n'),\n",
       " (1, 'num2'),\n",
       " (54, '='),\n",
       " (2, '6.3'),\n",
       " (4, '\\n'),\n",
       " (1, 'sum'),\n",
       " (54, '='),\n",
       " (1, 'var_1'),\n",
       " (54, '+'),\n",
       " (1, 'num2'),\n",
       " (4, '\\n'),\n",
       " (1, 'print'),\n",
       " (54, '('),\n",
       " (3, \"f'Sum: {sum}'\"),\n",
       " (54, ')'),\n",
       " (4, ''),\n",
       " (0, '')]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[0].Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "61e20fc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Input': [' ', 'write', 'a', 'python', 'function', 'to', 'add', 'two', 'user', 'provided', 'numbers', 'and', 'return', 'the', 'sum'], 'Output': [(63, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (54, '('), (1, 'num1'), (54, ','), (1, 'var_1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '    '), (1, 'sum'), (54, '='), (1, 'num1'), (54, '+'), (1, 'var_1'), (4, '\\n'), (1, 'return'), (1, 'sum'), (4, ''), (6, ''), (0, '')]}\n"
     ]
    }
   ],
   "source": [
    "print(vars(train_data.examples[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "18d24e45",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, \n",
    "                 input_dim, \n",
    "                 hid_dim, \n",
    "                 n_layers, \n",
    "                 n_heads, \n",
    "                 pf_dim,\n",
    "                 dropout, \n",
    "                 device,\n",
    "                 max_length = 1000):\n",
    "        super().__init__()\n",
    "\n",
    "        self.device = device\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(input_dim, hid_dim)\n",
    "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
    "        \n",
    "        self.layers = nn.ModuleList([EncoderLayer(hid_dim, \n",
    "                                                  n_heads, \n",
    "                                                  pf_dim,\n",
    "                                                  dropout, \n",
    "                                                  device) \n",
    "                                     for _ in range(n_layers)])\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n",
    "        \n",
    "    def forward(self, src, src_mask):\n",
    "        \n",
    "        #src = [batch size, src len]\n",
    "        #src_mask = [batch size, 1, 1, src len]\n",
    "        \n",
    "        batch_size = src.shape[0]\n",
    "        src_len = src.shape[1]\n",
    "\n",
    "        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n",
    "        \n",
    "        #pos = [batch size, src len]\n",
    "        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos))\n",
    "        \n",
    "        #src = [batch size, src len, hid dim]\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            src = layer(src, src_mask)\n",
    "            \n",
    "        #src = [batch size, src len, hid dim]\n",
    "            \n",
    "        return src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c4411a2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, \n",
    "                 hid_dim, \n",
    "                 n_heads, \n",
    "                 pf_dim,  \n",
    "                 dropout, \n",
    "                 device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
    "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, \n",
    "                                                                     pf_dim, \n",
    "                                                                     dropout)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src, src_mask):\n",
    "        \n",
    "        #src = [batch size, src len, hid dim]\n",
    "        #src_mask = [batch size, 1, 1, src len] \n",
    "                \n",
    "        #self attention\n",
    "        _src, _ = self.self_attention(src, src, src, src_mask)\n",
    "        \n",
    "        #dropout, residual connection and layer norm\n",
    "        src = self.self_attn_layer_norm(src + self.dropout(_src))\n",
    "        \n",
    "        #src = [batch size, src len, hid dim]\n",
    "        \n",
    "        #positionwise feedforward\n",
    "        _src = self.positionwise_feedforward(src)\n",
    "        \n",
    "        #dropout, residual and layer norm\n",
    "        src = self.ff_layer_norm(src + self.dropout(_src))\n",
    "        \n",
    "        #src = [batch size, src len, hid dim]\n",
    "        \n",
    "        return src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8916251b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionwiseFeedforwardLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, pf_dim, dropout):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.fc_1 = nn.Linear(hid_dim, pf_dim)\n",
    "        self.fc_2 = nn.Linear(pf_dim, hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        #x = [batch size, seq len, hid dim]\n",
    "        \n",
    "        x = self.dropout(torch.relu(self.fc_1(x)))\n",
    "        \n",
    "        #x = [batch size, seq len, pf dim]\n",
    "        \n",
    "        x = self.fc_2(x)\n",
    "        \n",
    "        #x = [batch size, seq len, hid dim]\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "f178aa3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, dropout, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        assert hid_dim % n_heads == 0\n",
    "        \n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_heads = n_heads\n",
    "        self.head_dim = hid_dim // n_heads\n",
    "        \n",
    "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_k = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_v = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_o = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n",
    "        \n",
    "    def forward(self, query, key, value, mask = None):\n",
    "        \n",
    "        batch_size = query.shape[0]\n",
    "        \n",
    "        #query = [batch size, query len, hid dim]\n",
    "        #key = [batch size, key len, hid dim]\n",
    "        #value = [batch size, value len, hid dim]\n",
    "                \n",
    "        Q = self.fc_q(query)\n",
    "        K = self.fc_k(key)\n",
    "        V = self.fc_v(value)\n",
    "        \n",
    "        #Q = [batch size, query len, hid dim]\n",
    "        #K = [batch size, key len, hid dim]\n",
    "        #V = [batch size, value len, hid dim]\n",
    "                \n",
    "        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        \n",
    "        #Q = [batch size, n heads, query len, head dim]\n",
    "        #K = [batch size, n heads, key len, head dim]\n",
    "        #V = [batch size, n heads, value len, head dim]\n",
    "                \n",
    "        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale\n",
    "        \n",
    "        #energy = [batch size, n heads, query len, key len]\n",
    "        \n",
    "        if mask is not None:\n",
    "            energy = energy.masked_fill(mask == 0, -1e10)\n",
    "        \n",
    "        attention = torch.softmax(energy, dim = -1)\n",
    "                \n",
    "        #attention = [batch size, n heads, query len, key len]\n",
    "                \n",
    "        x = torch.matmul(self.dropout(attention), V)\n",
    "        \n",
    "        #x = [batch size, n heads, query len, head dim]\n",
    "        \n",
    "        x = x.permute(0, 2, 1, 3).contiguous()\n",
    "        \n",
    "        #x = [batch size, query len, n heads, head dim]\n",
    "        \n",
    "        x = x.view(batch_size, -1, self.hid_dim)\n",
    "        \n",
    "        #x = [batch size, query len, hid dim]\n",
    "        \n",
    "        x = self.fc_o(x)\n",
    "        \n",
    "        #x = [batch size, query len, hid dim]\n",
    "        \n",
    "        return x, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "2e869229",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, \n",
    "                 output_dim, \n",
    "                 hid_dim, \n",
    "                 n_layers, \n",
    "                 n_heads, \n",
    "                 pf_dim, \n",
    "                 dropout, \n",
    "                 device,\n",
    "                 max_length = 10000):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.device = device\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(output_dim, hid_dim)\n",
    "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
    "        \n",
    "        self.layers = nn.ModuleList([DecoderLayer(hid_dim, \n",
    "                                                  n_heads, \n",
    "                                                  pf_dim, \n",
    "                                                  dropout, \n",
    "                                                  device)\n",
    "                                     for _ in range(n_layers)])\n",
    "        \n",
    "        self.fc_out = nn.Linear(hid_dim, output_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n",
    "        \n",
    "    def forward(self, trg, enc_src, trg_mask, src_mask):\n",
    "        \n",
    "        #trg = [batch size, trg len]\n",
    "        #enc_src = [batch size, src len, hid dim]\n",
    "        #trg_mask = [batch size, 1, trg len, trg len]\n",
    "        #src_mask = [batch size, 1, 1, src len]\n",
    "                \n",
    "        batch_size = trg.shape[0]\n",
    "        trg_len = trg.shape[1]\n",
    "        \n",
    "        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n",
    "                            \n",
    "        #pos = [batch size, trg len]\n",
    "\n",
    "        trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))\n",
    "                \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            trg, attention = layer(trg, enc_src, trg_mask, src_mask)\n",
    "        \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "        #attention = [batch size, n heads, trg len, src len]\n",
    "        \n",
    "        output = self.fc_out(trg)\n",
    "        \n",
    "        #output = [batch size, trg len, output dim]\n",
    "            \n",
    "        return output, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "0b44afb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, \n",
    "                 hid_dim, \n",
    "                 n_heads, \n",
    "                 pf_dim, \n",
    "                 dropout, \n",
    "                 device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.enc_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
    "        self.encoder_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
    "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, \n",
    "                                                                     pf_dim, \n",
    "                                                                     dropout)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, trg, enc_src, trg_mask, src_mask):\n",
    "        \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "        #enc_src = [batch size, src len, hid dim]\n",
    "        #trg_mask = [batch size, 1, trg len, trg len]\n",
    "        #src_mask = [batch size, 1, 1, src len]\n",
    "        \n",
    "        #self attention\n",
    "        _trg, _ = self.self_attention(trg, trg, trg, trg_mask)\n",
    "        \n",
    "        #dropout, residual connection and layer norm\n",
    "        trg = self.self_attn_layer_norm(trg + self.dropout(_trg))\n",
    "            \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "            \n",
    "        #encoder attention\n",
    "        _trg, attention = self.encoder_attention(trg, enc_src, enc_src, src_mask)\n",
    "        # query, key, value\n",
    "        \n",
    "        #dropout, residual connection and layer norm\n",
    "        trg = self.enc_attn_layer_norm(trg + self.dropout(_trg))\n",
    "                    \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "        \n",
    "        #positionwise feedforward\n",
    "        _trg = self.positionwise_feedforward(trg)\n",
    "        \n",
    "        #dropout, residual and layer norm\n",
    "        trg = self.ff_layer_norm(trg + self.dropout(_trg))\n",
    "        \n",
    "        #trg = [batch size, trg len, hid dim]\n",
    "        #attention = [batch size, n heads, trg len, src len]\n",
    "        \n",
    "        return trg, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "3dc5323a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, \n",
    "                 encoder, \n",
    "                 decoder, \n",
    "                 src_pad_idx, \n",
    "                 trg_pad_idx, \n",
    "                 device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "        self.device = device\n",
    "        \n",
    "    def make_src_mask(self, src):\n",
    "        \n",
    "        #src = [batch size, src len]\n",
    "        \n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
    "\n",
    "        #src_mask = [batch size, 1, 1, src len]\n",
    "\n",
    "        return src_mask\n",
    "    \n",
    "    def make_trg_mask(self, trg):\n",
    "        \n",
    "        #trg = [batch size, trg len]\n",
    "        \n",
    "        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)\n",
    "        \n",
    "        #trg_pad_mask = [batch size, 1, 1, trg len]\n",
    "        \n",
    "        trg_len = trg.shape[1]\n",
    "        \n",
    "        trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device = self.device)).bool()\n",
    "        \n",
    "        #trg_sub_mask = [trg len, trg len]\n",
    "            \n",
    "        trg_mask = trg_pad_mask & trg_sub_mask\n",
    "        \n",
    "        #trg_mask = [batch size, 1, trg len, trg len]\n",
    "        \n",
    "        return trg_mask\n",
    "\n",
    "    def forward(self, src, trg):\n",
    "        \n",
    "        #src = [batch size, src len]\n",
    "        #trg = [batch size, trg len]\n",
    "                \n",
    "        src_mask = self.make_src_mask(src)\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        \n",
    "        #src_mask = [batch size, 1, 1, src len]\n",
    "        #trg_mask = [batch size, 1, trg len, trg len]\n",
    "        \n",
    "        enc_src = self.encoder(src, src_mask)\n",
    "        \n",
    "        #enc_src = [batch size, src len, hid dim]\n",
    "                \n",
    "        output, attention = self.decoder(trg, enc_src, trg_mask, src_mask)\n",
    "        \n",
    "        #output = [batch size, trg len, output dim]\n",
    "        #attention = [batch size, n heads, trg len, src len]\n",
    "        \n",
    "        return output, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "dba4499d",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = len(Input.vocab)\n",
    "OUTPUT_DIM = len(Output.vocab)\n",
    "HID_DIM = 256\n",
    "ENC_LAYERS = 3\n",
    "DEC_LAYERS = 3\n",
    "ENC_HEADS = 16\n",
    "DEC_HEADS = 16\n",
    "ENC_PF_DIM = 512\n",
    "DEC_PF_DIM = 512\n",
    "ENC_DROPOUT = 0.1\n",
    "DEC_DROPOUT = 0.1\n",
    "\n",
    "enc = Encoder(INPUT_DIM, \n",
    "              HID_DIM, \n",
    "              ENC_LAYERS, \n",
    "              ENC_HEADS, \n",
    "              ENC_PF_DIM, \n",
    "              ENC_DROPOUT, \n",
    "              device)\n",
    "\n",
    "dec = Decoder(OUTPUT_DIM, \n",
    "              HID_DIM, \n",
    "              DEC_LAYERS, \n",
    "              DEC_HEADS, \n",
    "              DEC_PF_DIM, \n",
    "              DEC_DROPOUT, \n",
    "              device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "fc49c866",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5668"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(Output.vocab.__dict__['freqs'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e630e18d",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC_PAD_IDX = Input.vocab.stoi[Input.pad_token]\n",
    "TRG_PAD_IDX = Output.vocab.stoi[Output.pad_token]\n",
    "\n",
    "model = Seq2Seq(enc, dec, SRC_PAD_IDX, TRG_PAD_IDX, device).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "fe6e56e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 10,220,071 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "9d5801e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_weights(m):\n",
    "    if hasattr(m, 'weight') and m.weight.dim() > 1:\n",
    "        nn.init.xavier_uniform_(m.weight.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "bd7233f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.apply(initialize_weights);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "55984f5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 0.0005\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fb444f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import math\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class CrossEntropyLoss(nn.CrossEntropyLoss):\n",
    "    \"\"\"CrossEntropyLoss - with ability to recieve distrbution as targets, and optional label smoothing\"\"\"\n",
    "\n",
    "    def __init__(self, weight=None, ignore_index=-100, reduction='mean', smooth_eps=None, smooth_dist=None, from_logits=True):\n",
    "        super(CrossEntropyLoss, self).__init__(weight=weight,\n",
    "                                               ignore_index=ignore_index, reduction=reduction)\n",
    "        self.smooth_eps = smooth_eps\n",
    "        self.smooth_dist = smooth_dist\n",
    "        self.from_logits = from_logits\n",
    "\n",
    "    def forward(self, input, target, smooth_dist=None):\n",
    "        if smooth_dist is None:\n",
    "            smooth_dist = self.smooth_dist\n",
    "        return cross_entropy(input, target, weight=self.weight, ignore_index=self.ignore_index,\n",
    "                             reduction=self.reduction, smooth_eps=self.smooth_eps,\n",
    "                             smooth_dist=smooth_dist, from_logits=self.from_logits)\n",
    "\n",
    "\n",
    "def cross_entropy(inputs, target, weight=None, ignore_index=-100, reduction='mean',\n",
    "                  smooth_eps=None, smooth_dist=None, from_logits=True):\n",
    "    \"\"\"cross entropy loss, with support for target distributions and label smoothing https://arxiv.org/abs/1512.00567\"\"\"\n",
    "    smooth_eps = smooth_eps or 0\n",
    "\n",
    "    # ordinary log-liklihood - use cross_entropy from nn\n",
    "    if _is_long(target) and smooth_eps == 0:\n",
    "        if from_logits:\n",
    "            return F.cross_entropy(inputs, target, weight, ignore_index=ignore_index, reduction=reduction)\n",
    "        else:\n",
    "            return F.nll_loss(inputs, target, weight, ignore_index=ignore_index, reduction=reduction)\n",
    "\n",
    "    if from_logits:\n",
    "        # log-softmax of inputs\n",
    "        lsm = F.log_softmax(inputs, dim=-1)\n",
    "    else:\n",
    "        lsm = inputs\n",
    "\n",
    "    masked_indices = None\n",
    "    num_classes = inputs.size(-1)\n",
    "\n",
    "    if _is_long(target) and ignore_index >= 0:\n",
    "        masked_indices = target.eq(ignore_index)\n",
    "\n",
    "    if smooth_eps > 0 and smooth_dist is not None:\n",
    "        if _is_long(target):\n",
    "            target = onehot(target, num_classes).type_as(inputs)\n",
    "        if smooth_dist.dim() < target.dim():\n",
    "            smooth_dist = smooth_dist.unsqueeze(0)\n",
    "        target.lerp_(smooth_dist, smooth_eps)\n",
    "\n",
    "    if weight is not None:\n",
    "        lsm = lsm * weight.unsqueeze(0)\n",
    "\n",
    "    if _is_long(target):\n",
    "        eps_sum = smooth_eps / num_classes\n",
    "        eps_nll = 1. - eps_sum - smooth_eps\n",
    "        likelihood = lsm.gather(dim=-1, index=target.unsqueeze(-1)).squeeze(-1)\n",
    "        loss = -(eps_nll * likelihood + eps_sum * lsm.sum(-1))\n",
    "    else:\n",
    "        loss = -(target * lsm).sum(-1)\n",
    "\n",
    "    if masked_indices is not None:\n",
    "        loss.masked_fill_(masked_indices, 0)\n",
    "\n",
    "    if reduction == 'sum':\n",
    "        loss = loss.sum()\n",
    "    elif reduction == 'mean':\n",
    "        if masked_indices is None:\n",
    "            loss = loss.mean()\n",
    "        else:\n",
    "            loss = loss.sum() / float(loss.size(0) - masked_indices.sum())\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def onehot(indexes, N=None, ignore_index=None):\n",
    "    \"\"\"\n",
    "    Creates a one-representation of indexes with N possible entries\n",
    "    if N is not specified, it will suit the maximum index appearing.\n",
    "    indexes is a long-tensor of indexes\n",
    "    ignore_index will be zero in onehot representation\n",
    "    \"\"\"\n",
    "    if N is None:\n",
    "        N = indexes.max() + 1\n",
    "    sz = list(indexes.size())\n",
    "    output = indexes.new().byte().resize_(*sz, N).zero_()\n",
    "    output.scatter_(-1, indexes.unsqueeze(-1), 1)\n",
    "    if ignore_index is not None and ignore_index >= 0:\n",
    "        output.masked_fill_(indexes.eq(ignore_index).unsqueeze(-1), 0)\n",
    "    return output\n",
    "\n",
    "def _is_long(x):\n",
    "    if hasattr(x, 'data'):\n",
    "        x = x.data\n",
    "    return isinstance(x, torch.LongTensor) or isinstance(x, torch.cuda.LongTensor)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c01b3bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def maskNLLLoss(inp, target, mask):\n",
    "    # print(inp.shape, target.shape, mask.sum())\n",
    "    nTotal = mask.sum()\n",
    "    crossEntropy = CrossEntropyLoss(ignore_index = TRG_PAD_IDX, smooth_eps=0.20)\n",
    "    loss = crossEntropy(inp, target)\n",
    "    loss = loss.to(device)\n",
    "    return loss, nTotal.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "f9bba37d",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = maskNLLLoss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "8c3bb883",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def make_trg_mask(trg):\n",
    "        \n",
    "        #trg = [batch size, trg len]\n",
    "        \n",
    "        trg_pad_mask = (trg != TRG_PAD_IDX).unsqueeze(1).unsqueeze(2)\n",
    "        \n",
    "        #trg_pad_mask = [batch size, 1, 1, trg len]\n",
    "        \n",
    "        trg_len = trg.shape[1]\n",
    "        \n",
    "        trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device = device)).bool()\n",
    "        \n",
    "        #trg_sub_mask = [trg len, trg len]\n",
    "            \n",
    "        trg_mask = trg_pad_mask & trg_sub_mask\n",
    "        \n",
    "        #trg_mask = [batch size, 1, trg len, trg len]\n",
    "        \n",
    "        return trg_mask\n",
    "\n",
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    n_totals = 0\n",
    "    print_losses = []\n",
    "    for i, batch in tqdm(enumerate(iterator), total=len(iterator)):\n",
    "        # print(batch)\n",
    "        loss = 0\n",
    "        src = batch.Input.permute(1, 0)\n",
    "        trg = batch.Output.permute(1, 0)\n",
    "        trg_mask = make_trg_mask(trg)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output, _ = model(src, trg[:,:-1])\n",
    "                \n",
    "        #output = [batch size, trg len - 1, output dim]\n",
    "        #trg = [batch size, trg len]\n",
    "            \n",
    "        output_dim = output.shape[-1]\n",
    "            \n",
    "        output = output.contiguous().view(-1, output_dim)\n",
    "        trg = trg[:,1:].contiguous().view(-1)\n",
    "                \n",
    "        #output = [batch size * trg len - 1, output dim]\n",
    "        #trg = [batch size * trg len - 1]\n",
    "            \n",
    "        mask_loss, nTotal = criterion(output, trg, trg_mask)\n",
    "        \n",
    "        mask_loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        print_losses.append(mask_loss.item() * nTotal)\n",
    "        n_totals += nTotal\n",
    "\n",
    "\n",
    "        \n",
    "    return sum(print_losses) / n_totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "8189ed56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    n_totals = 0\n",
    "    print_losses = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in tqdm(enumerate(iterator), total=len(iterator)):\n",
    "\n",
    "            src = batch.Input.permute(1, 0)\n",
    "            trg = batch.Output.permute(1, 0)\n",
    "            trg_mask = make_trg_mask(trg)\n",
    "\n",
    "            output, _ = model(src, trg[:,:-1])\n",
    "            \n",
    "            #output = [batch size, trg len - 1, output dim]\n",
    "            #trg = [batch size, trg len]\n",
    "            \n",
    "            output_dim = output.shape[-1]\n",
    "            \n",
    "            output = output.contiguous().view(-1, output_dim)\n",
    "            trg = trg[:,1:].contiguous().view(-1)\n",
    "            \n",
    "            #output = [batch size * trg len - 1, output dim]\n",
    "            #trg = [batch size * trg len - 1]\n",
    "            \n",
    "            mask_loss, nTotal = criterion(output, trg, trg_mask)\n",
    "\n",
    "            print_losses.append(mask_loss.item() * nTotal)\n",
    "            n_totals += nTotal\n",
    "\n",
    "        \n",
    "    return sum(print_losses) / n_totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "029b9987",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def epoch_time(start_time, end_time):\n",
    "#     elapsed_time = end_time - start_time\n",
    "#     elapsed_mins = int(elapsed_time / 60)\n",
    "#     elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "#     return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "26ea9777",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# N_EPOCHS = 50\n",
    "# CLIP = 1\n",
    "\n",
    "# best_valid_loss = float('inf')\n",
    "\n",
    "# for epoch in range(N_EPOCHS):\n",
    "    \n",
    "#     start_time = time.time()\n",
    "    \n",
    "#     train_example = []\n",
    "#     val_example = []\n",
    "\n",
    "#     for i in range(train_df.shape[0]):\n",
    "#         try:\n",
    "#             ex = data.Example.fromlist([train_df.question[i], train_df.solution[i]], fields)\n",
    "#             train_example.append(ex)\n",
    "#         except:\n",
    "#             pass\n",
    "\n",
    "#     for i in range(val_df.shape[0]):\n",
    "#         try:\n",
    "#             ex = data.Example.fromlist([val_df.question[i], val_df.solution[i]], fields)\n",
    "#             val_example.append(ex)\n",
    "#         except:\n",
    "#             pass       \n",
    "\n",
    "#     train_data = data.Dataset(train_example, fields)\n",
    "#     valid_data =  data.Dataset(val_example, fields)\n",
    "\n",
    "#     BATCH_SIZE = 16\n",
    "#     train_iterator, valid_iterator = BucketIterator.splits((train_data, valid_data), batch_size = BATCH_SIZE, \n",
    "#                                                                 sort_key = lambda x: len(x.Input),\n",
    "#                                                                 sort_within_batch=True, device = device)\n",
    "\n",
    "#     train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
    "#     valid_loss = evaluate(model, valid_iterator, criterion)\n",
    "    \n",
    "#     end_time = time.time()\n",
    "    \n",
    "#     epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "#     if valid_loss < best_valid_loss:\n",
    "#         best_valid_loss = valid_loss\n",
    "#         torch.save(model.state_dict(), 'C:/Users/suche/OneDrive/Desktop/FYP/project/model1.pt')\n",
    "    \n",
    "#     print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "#     print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "#     print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "b2afa87a",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC = Input\n",
    "TRG = Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "ee98125f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('C:/Users/suche/OneDrive/Desktop/FYP/project/model1.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "782b11a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate_sentence(sentence, src_field, trg_field, model, device, max_len = 50000):\n",
    "    \n",
    "    model.eval()\n",
    "        \n",
    "    if isinstance(sentence, str):\n",
    "        nlp = spacy.load('en')\n",
    "        tokens = [token.text.lower() for token in nlp(sentence)]\n",
    "    else:\n",
    "        tokens = [token.lower() for token in sentence]\n",
    "\n",
    "    tokens = [src_field.init_token] + tokens + [src_field.eos_token]\n",
    "        \n",
    "    src_indexes = [src_field.vocab.stoi[token] for token in tokens]\n",
    "\n",
    "    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)\n",
    "    \n",
    "    src_mask = model.make_src_mask(src_tensor)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        enc_src = model.encoder(src_tensor, src_mask)\n",
    "\n",
    "    trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]\n",
    "\n",
    "    for i in range(max_len):\n",
    "\n",
    "        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)\n",
    "\n",
    "        trg_mask = model.make_trg_mask(trg_tensor)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            output, attention = model.decoder(trg_tensor, enc_src, trg_mask, src_mask)\n",
    "        \n",
    "        pred_token = output.argmax(2)[:,-1].item()\n",
    "        \n",
    "        trg_indexes.append(pred_token)\n",
    "\n",
    "        if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:\n",
    "            break\n",
    "    \n",
    "    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]\n",
    "    \n",
    "    return trg_tokens[1:], attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "ac1b66dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_attention(sentence, translation, attention, n_heads = 8, n_rows = 4, n_cols = 2):\n",
    "    \n",
    "    assert n_rows * n_cols == n_heads\n",
    "    \n",
    "    fig = plt.figure(figsize=(30,50))\n",
    "    \n",
    "    for i in range(n_heads):\n",
    "        \n",
    "        ax = fig.add_subplot(n_rows, n_cols, i+1)\n",
    "        \n",
    "        _attention = attention.squeeze(0)[i].cpu().detach().numpy()\n",
    "\n",
    "        cax = ax.matshow(_attention, cmap='bone')\n",
    "\n",
    "        ax.tick_params(labelsize=12)\n",
    "        ax.set_xticklabels(['']+['<sos>']+[t.lower() for t in sentence]+['<eos>'], \n",
    "                           rotation=45)\n",
    "        ax.set_yticklabels(['']+translation)\n",
    "\n",
    "        ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "        ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "f6d4c75c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg sequence: \n",
      "[(63, 'utf-8'), (1, 'def'), (1, 'gcd'), (54, '('), (1, 'var_1'), (54, ','), (1, 'y'), (54, ')'), (54, ':'), (4, '\\n'), (5, '   '), (1, 'if'), (1, 'var_1'), (54, '>'), (1, 'y'), (54, ':'), (4, '\\n'), (5, '      '), (1, 'var_2'), (54, '='), (1, 'var_1'), (4, '\\n'), (6, ''), (1, 'else'), (54, ':'), (4, '\\n'), (5, '      '), (1, 'var_2'), (54, '='), (1, 'var_1'), (4, '\\n'), (62, '\\n'), (6, ''), (1, 'for'), (1, 'i'), (1, 'in'), (1, 'range'), (54, '('), (2, '1'), (54, ','), (1, 'var_2'), (54, '+'), (2, '1'), (54, ')'), (54, ':'), (4, '\\n'), (5, '      '), (1, 'if'), (54, '('), (54, '('), (1, 'var_1'), (54, '%'), (1, 'i'), (54, '=='), (2, '0'), (54, ')'), (1, 'and'), (54, '('), (1, 'var_2'), (54, '%'), (1, 'i'), (54, '=='), (2, '0'), (54, ')'), (54, ')'), (54, ':'), (4, '\\n'), (5, '         '), (1, 'gcd'), (54, '='), (1, 'i'), (4, '\\n'), (62, '\\n'), (6, ''), (6, ''), (1, 'return'), (1, 'gcd'), (4, ''), (6, ''), (0, ''), '']\n",
      "code: \n",
      " def gcd (var_1 ,y ):\n",
      "   if var_1 >y :\n",
      "      var_2 =var_1 \n",
      "   else :\n",
      "      var_2 =var_1 \n",
      "\n",
      "   for i in range (1 ,var_2 +1 ):\n",
      "      if ((var_1 %i ==0 )and (var_2 %i ==0 )):\n",
      "         gcd =i \n",
      "\n",
      "   return gcd \n"
     ]
    }
   ],
   "source": [
    "src = \"write a function that adds two numbers\"\n",
    "src=src.split(\" \")\n",
    "translation, attention = translate_sentence(src, SRC, TRG, model, device)\n",
    "\n",
    "print(f'predicted trg sequence: ')\n",
    "print(translation)\n",
    "print(\"code: \\n\", untokenize(translation[:-1]).decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "3e555bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display_attention(src, translation, attention)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "9393d8b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('C:/Users/suche/OneDrive/Desktop/FYP/project/model1.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "614b7190",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eng_to_python(src):\n",
    "  src=src.split(\" \")\n",
    "  translation, attention = translate_sentence(src, SRC, TRG, model, device)\n",
    "\n",
    "  print(f'predicted trg: \\n')\n",
    "  # print(translation)\n",
    "  print(untokenize(translation[:-1]).decode('utf-8'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "5078b588",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC = Input\n",
    "TRG = Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "19aa27b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg: \n",
      "\n",
      "def merge_lists (l1 :list ,l2 :list ):\n",
      "    return list (zip (l1 ,l2 ))\n"
     ]
    }
   ],
   "source": [
    "src = \"function to merge two lists\"\n",
    "\n",
    "eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "1b895442",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg: \n",
      "\n",
      "print (\"Hello World\")\n"
     ]
    }
   ],
   "source": [
    "src = \"program to print hello world\"\n",
    "\n",
    "eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "da445a17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg: \n",
      "\n",
      "\n",
      "def multiply (numbers ):\n",
      "    var_1 =1 \n",
      "    for x in numbers :\n",
      "        var_1 *=x \n",
      "    return var_1 \n",
      "\n",
      "print (multiply (8 ,2 ,3 ,-1 ,7 ))\n"
     ]
    }
   ],
   "source": [
    "src = \"program to multiply integers in a list\"\n",
    "\n",
    "eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "099ccb27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg: \n",
      "\n",
      "\n",
      "def reverse_string (string ):\n",
      "    return string [::-1 ]\n"
     ]
    }
   ],
   "source": [
    "src = \"program to reverse a string\"\n",
    "\n",
    "eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "a1d2bb2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted trg: \n",
      "\n",
      "import re \n",
      "\n",
      "def removeLeadingZeros (var_1 ):\n",
      "    var_2 =re .sub (regex ,'.',var_1 )\n",
      "    print (var_2 )\n",
      "\n",
      "if __name__ =='__main__':\n",
      "\n",
      "    var_3 =\"216.08.094.196\"\n",
      "    removeLeadingZeros (var_2 )\n"
     ]
    }
   ],
   "source": [
    "src = \"write a program to extract the mobile number from the given string in Python \"\n",
    "\n",
    "eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "3b6283b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.load_state_dict(torch.load('C:/Users/suche/OneDrive/Desktop/FYP/project/model1.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "ff62754a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Define the model architecture\n",
    "# model = Seq2Seq(input_dim, output_dim, enc_hid_dim, dec_hid_dim, dropout)\n",
    "# # Load the saved parameters\n",
    "# model.load_state_dict(torch.load('C:/Users/suche/OneDrive/Desktop/FYP/project/model1.pt'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "69372a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# src = \"program to add two numbers\"\n",
    "\n",
    "# eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "c5abb1db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# src = \"program to check whether number is palindrome or not.\"\n",
    "\n",
    "# eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "e3d06716",
   "metadata": {},
   "outputs": [],
   "source": [
    "# src = \"program to count the characters in a sentence\"\n",
    "\n",
    "# eng_to_python(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "d461ccf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# src = \"program to print HELLO\"\n",
    "\n",
    "# eng_to_python(src)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test",
   "language": "python",
   "name": "test"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
